import pandas as pd
import requests
import time
import re
from urllib.parse import quote
from tqdm import tqdm

def extract_cas_from_title(title):
    # Pattern for CAS number: nums-nums-num
    cas_pattern = r'\b\d+-\d+-\d+\b'
    match = re.search(cas_pattern, title)
    return match.group(0) if match else ''

def search_pubchem(compound_name):

    try:
        # URL encode the compound name
        encoded_name = quote(compound_name)
        
        # CID
        search_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{encoded_name}/cids/JSON"
        response = requests.get(search_url)
        
        if response.status_code == 200:
            cid = str(response.json()['IdentifierList']['CID'][0])  # Get the best match
            
            # Get the title
            summary_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/synonyms/JSON"
            summary_response = requests.get(summary_url)
            
            cas = ''
            if summary_response.status_code == 200:
                synonyms = summary_response.json()['InformationList']['Information'][0]['Synonym']
                # CAS 
                for synonym in synonyms:
                    cas_match = extract_cas_from_title(synonym)
                    if cas_match:
                        cas = cas_match
                        break
            
            # SMILES and Molecular Formula
            property_url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/property/CanonicalSMILES,MolecularFormula/JSON"
            prop_response = requests.get(property_url)
            
            if prop_response.status_code == 200:
                properties = prop_response.json()['PropertyTable']['Properties'][0]
                smiles = properties['CanonicalSMILES']
                molecular_formula = properties['MolecularFormula']
                
                return {
                    'CID': cid,
                    'SMILES': smiles,
                    'CAS': cas,
                    'MolecularFormula': molecular_formula
                }
            
        return None
    except Exception as e:
        print(f"Error searching PubChem for {compound_name}: {str(e)}")
        return None
    

def process_excel_file(excel_path):
    # Generate a df
    df = pd.read_excel(excel_path)
    df['CID'] = ''
    df['SMILES'] = ''
    df['CAS'] = ''
    df['MolecularFormula'] = ''
    
    # Process each compound
    for idx, row in tqdm(df.iterrows(), total=len(df)):
        compound_name = row['Compound Name']
        if pd.notna(compound_name):  # Check if compound name is not NaN
            result = search_pubchem(compound_name)
            if result:
                df.at[idx, 'CID'] = result['CID']
                df.at[idx, 'SMILES'] = result['SMILES']
                df.at[idx, 'CAS'] = result['CAS']
                df.at[idx, 'MolecularFormula'] = result['MolecularFormula']
            
            # Add a delay to avoid overwhelming the API
            time.sleep(0.2)
    
    return df

# Usage
results_df = process_excel_file("compound_names.xlsx")

results_df.to_excel("results_with_formula.xlsx", index=False)